package sc.nio;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.channels.spi.SelectorProvider;
import java.util.*;

import sc.server.Server;

/**
 * Threaded socket server, that uses a Selector to multiplex SocketChannels.
 *
 * @author stephencarmody@gmail.com Stephen Carmody
 */
public class NioServer implements Runnable
{
	// The host:port combination to listen on
	private InetAddress hostAddress;
	private int port;

	// The channel on which we'll accept connections
	private ServerSocketChannel serverChannel;

	// The selector we'll be monitoring
	private Selector selector;

	// A list of PendingChange instances
	private List<NioChangeRequest> pendingChanges = new ArrayList<NioChangeRequest>();

	// Maps a SocketChannel to a list of ByteBuffer instances
	private Map<SocketChannel, List<ByteBuffer>> pendingData = new HashMap<SocketChannel, List<ByteBuffer>>();

	// The handler for the received client messages
	private NioHandler handler;

	// Buffer to read data off a channel into
	private ByteBuffer readBuffer = ByteBuffer.allocate(8192);

	// Flags logging of reads and writes
	private boolean logIO;
	
	/**
	 * Creates an instance of NioServer.
	 * 
	 * @param hostAddress
	 *            The host address the server is running on
	 * @param port
	 *            The port to listen on
	 * @param handler
	 *            The NioReader instance to handle received channel data
	 * 
	 * @throws IOException
	 */
	public NioServer(InetAddress hostAddress, int port, NioHandler handler)
			throws IOException
	{
		this.hostAddress = hostAddress;
		this.port = port;

		// Initialise the Selector
		this.selector = SelectorProvider.provider().openSelector();
		this.serverChannel = ServerSocketChannel.open();
		this.serverChannel.configureBlocking(false);
		InetSocketAddress isa = new InetSocketAddress(this.hostAddress,
				this.port);
		this.serverChannel.socket().bind(isa);
		this.serverChannel.register(this.selector, SelectionKey.OP_ACCEPT);

		this.handler = handler;
	}

	/**
	 * Switches the logging of IO on/off.
	 * 
	 * @param logIO
	 * 			true for logging, false otherwise
	 */
	public void setLogIO(boolean logIO)
	{
		this.logIO = logIO;
	}
	
	/**
	 * Sends the specified data through the specified channel.
	 * 
	 * @param channel
	 *            A SocketChannel instance
	 * @param data
	 *            The data to be sent *
	 */
	public void send(SocketChannel channel, byte[] data)
	{
		synchronized (this.pendingChanges)
		{
			// Queue interest in writing to the channel
			this.pendingChanges.add(new NioChangeRequest(channel,
					SelectionKey.OP_WRITE));

			// And queue the data we want written
			synchronized (this.pendingData)
			{
				List<ByteBuffer> queue = this.pendingData.get(channel);
				if (queue == null)
				{
					queue = new ArrayList<ByteBuffer>();
					this.pendingData.put(channel, queue);
				}
				queue.add(ByteBuffer.wrap(data));
			}
		}

		// Wake up our selecting thread so it can make the required changes
		this.selector.wakeup();
	}

	/**
	 * Thread main loop.
	 */
	public void run()
	{
		while (true)
		{
			try
			{
				// Process any pending changes
				synchronized (this.pendingChanges)
				{
					for (NioChangeRequest change : this.pendingChanges)
					{
						SelectionKey key = change.socket.keyFor(this.selector);
						key.interestOps(change.ops);
					}
					this.pendingChanges.clear();
				}

				// Wait for an event one of the registered channels
				this.selector.select();

				// Iterate over the set of keys for which events are available
				Set<SelectionKey> selectedkeys = this.selector.selectedKeys();
				for (SelectionKey key : selectedkeys)
				{
					selectedkeys.remove(key);

					// Check what event is available and deal with it
					if (key.isAcceptable())
						this.accept(key);
					else if (key.isReadable())
						this.read(key);
					else if (key.isWritable())
						this.write(key);
				}
			}
			catch (Exception e)
			{
				e.printStackTrace();
			}
		}
	}

	/**
	 * Closes the specified SocketChannel and cancels the specified
	 * SelectionKey.
	 * 
	 * @param channel
	 *            A SocketChannel instance
	 * @param key
	 *            A SelectionKey instance
	 * 
	 * @throws IOException
	 */
	private void closeChannel(SocketChannel channel, SelectionKey key)
			throws IOException
	{
		key.cancel();
		channel.close();
		this.pendingData.remove(channel);
	}

	/**
	 * Called when a channel connection is initiated and needs to be accepted.
	 * 
	 * @param key
	 *            The SelectionKey of channel
	 * 
	 * @throws IOException
	 */
	private void accept(SelectionKey key) throws IOException
	{
		// For an accept to be pending the channel must be a server socket
		// channel.
		ServerSocketChannel serverSocketChannel = (ServerSocketChannel) key
				.channel();

		// Accept the connection and make it non-blocking
		SocketChannel socketChannel = serverSocketChannel.accept();
		socketChannel.configureBlocking(false);

		// Register new SocketChannel with our Selector, indicating
		// we'd like to be notified when there's data waiting to be read
		socketChannel.register(this.selector, SelectionKey.OP_READ);
	}

	/**
	 * Called when a connected channel has data waiting to be read.
	 * 
	 * @param key
	 *            The SelectionKey of channel
	 * 
	 * @throws IOException
	 */
	private void read(SelectionKey key) throws IOException
	{
		int numRead;
		SocketChannel channel = (SocketChannel) key.channel();

		try
		{
			// Attempt to read off the channel
			numRead = channel.read(this.readBuffer);
		}
		catch (IOException e)
		{
			// The client forcibly closed the connection
			this.closeChannel(channel, key);
			return;
		}

		if (numRead == -1)
		{
			// The client shut the socket down cleanly
			this.closeChannel(channel, key);
			return;
		}

		// Hand a copy of the data off to our reading thread
		ByteBuffer buffer = ByteBuffer.allocate(numRead);
		buffer.put(this.readBuffer);
		if (this.logIO)
			Server.logger.info("NIO <- " + buffer.toString());
		NioChannelData channeldata = new NioChannelData(channel, buffer);
		this.handler.receive(channeldata);
		this.readBuffer.clear();
	}

	/**
	 * Called when a connected channel has data waiting to be written.
	 * 
	 * @param key
	 *            The SelectionKey of channel
	 * 
	 * @throws IOException
	 */
	private void write(SelectionKey key) throws IOException
	{
		SocketChannel socketChannel = (SocketChannel) key.channel();

		synchronized (this.pendingData)
		{
			List<ByteBuffer> queue = this.pendingData.get(socketChannel);

			// Write until there's not more data ...
			while (!queue.isEmpty())
			{
				ByteBuffer buf = (ByteBuffer) queue.get(0);
				if (this.logIO)
					Server.logger.info("NIO -> " + Arrays.toString(buf.array()) + "\"" + new String(buf.array()) + "\"");
				socketChannel.write(buf);
				// ... or the socket's buffer fills up
				if (buf.remaining() > 0)
					break;
				queue.remove(0);
			}

			// We wrote away all data, so we're no longer interested
			// in writing on this socket. Switch back to waiting for
			// data.
			if (queue.isEmpty())
				key.interestOps(SelectionKey.OP_READ);
		}
	}
}
